#include "FileReader.hpp"
#include "H264CommonFunction.hpp"


RBSP::RBSP(){}

RBSP::~RBSP(){
    if(RBSPBuf != nullptr){
        free(RBSPBuf);
        RBSPBuf = nullptr;
    }
}


EBSP::EBSP(){}

EBSP::~EBSP(){
    if(EBSPBuf != nullptr){
        free(EBSPBuf);
        EBSPBuf = nullptr;
    }
}


int EBSP::GetRBSP(RBSP& rbsp){
    int index = 0;
    uint8_t* startBuf = EBSPBuf;
    rbsp.RBSPLen = EBSPLen;
    rbsp.RBSPBuf = (uint8_t*)malloc(rbsp.RBSPLen);
    for(int i = 0; i < EBSPLen - 4; i++){
        if(startBuf[i] == 0 && startBuf[i+1] == 0 && startBuf[i+2] == 3 && (startBuf[i+3] == 0 ||startBuf[i+3] == 1
       || startBuf[i+3] == 2 || startBuf[i+3] == 3)){
            --rbsp.RBSPLen;
            continue;
       }
       rbsp.RBSPBuf[index] = startBuf[i];
       index++;
    }
    return 0;
}


Nalu::Nalu(){}

Nalu::~Nalu(){
    if(buf != nullptr){
        free(buf);
        buf = nullptr;
    }
}

int Nalu::SetBuf(uint8_t* bufPtr, int len){
    if(buf != nullptr){
        free(buf);
        buf = nullptr;
    }
    bufLen = len;

    buf = (uint8_t*)malloc(len);
    memcpy(buf, bufPtr, len);
    return 0;
}

int Nalu::GetEBSP(EBSP& ebsp){
    ebsp.EBSPLen = Nalu::bufLen - Nalu::StartCodeLen;

    ebsp.EBSPBuf = (uint8_t*)malloc(ebsp.EBSPLen);
    memcpy(ebsp.EBSPBuf, Nalu::buf + Nalu::StartCodeLen, ebsp.EBSPLen);

    return 0;
}

int Nalu::ParseRBSP(){
    EBSP esbp;
    int ret = GetEBSP(esbp);
    if(ret != 0){
        return -1;
    }
    return esbp.GetRBSP(rbsp);
}

int Nalu::ParseHeader(){
    uint8_t header = rbsp.RBSPBuf[0];
    forbidden_zero_bit = (header>>7) & 0x1;
    nal_ref_idc = (header >> 5)& 0x3;
    nal_unit_type = header & 0x1f;
    IdrPicFlag = (nal_unit_type == H264_NAL_IDR_SLICE) ? 1 : 0;
    return 0;
}



AnnexBReader::AnnexBReader(std::string& _fileName){
    fileName = (char*)_fileName.c_str();
}

AnnexBReader::~AnnexBReader(){
    if(buffer != nullptr){
        free(buffer);
        buffer = nullptr;
    }
    if(f != nullptr){
        fclose(f);
        f = nullptr;
    }
}

int AnnexBReader::Open(){
    f = fopen(fileName, "rb");
    if(f == nullptr){
        return -1;
    }
    return 0;
}

int AnnexBReader::Close(){
    if(f == nullptr){
        fclose(f);
        f = nullptr;
    }
    if(buffer != nullptr){
        free(buffer);
        buffer = nullptr;
    }
    return 0;
}

int AnnexBReader::ReadFromFile(){
    int tempBufferLen = 1024;
    uint8_t* buf = (uint8_t*)malloc(tempBufferLen);
    int readedLen = fread(buf, 1, tempBufferLen, f);

    if(readedLen > 0){
        uint8_t* _buf = (uint8_t*)malloc(tempBufferLen + bufferLen);
        memcpy(_buf, buffer, bufferLen);
        memcpy(_buf+bufferLen, buf, tempBufferLen );
        bufferLen += tempBufferLen;

        if(buffer != nullptr){
            free(buffer);
            buffer = nullptr;
        }

        buffer = _buf;
    }

    free(buf);
    return readedLen;
}

bool AnnexBReader::CheckStartCode(int& StartCodeLen, uint8_t* bufPtr, int bufLen){
    if(bufLen <= 2){
        StartCodeLen = 0;
        return false;
    }
    else if(bufLen >= 4){
        if(bufPtr[0] == 0){
            if(bufPtr[1] == 0){
                if(bufPtr[2] == 0){
                    if(bufPtr[3] == 1){
                        StartCodeLen = 4;
                        return true;
                    }
                }
                else if(bufPtr[2] == 1){
                    StartCodeLen = 3;
                    return true;
                }
            }
        }
    }
    else if(bufLen <= 3){
       if(bufPtr[0] == 0){
            if(bufPtr[1] == 0){
                if(bufPtr[2] == 1){
                    StartCodeLen = 3;
                    return true;
                }
            }
       }
    }
    StartCodeLen = 0;
    return false;

}


int AnnexBReader::ReadNalu(Nalu& nalu){
    //printf("=====================111=\n");

    while(1){
        if(bufferLen <= 0){
            int readedLen = ReadFromFile();
            if(readedLen < 0){
                isEnd = true;
            }
        }

        uint8_t* buf = buffer;

        int startCodeLen = 0;
        bool isStartCode = CheckStartCode(startCodeLen, buf, bufferLen);
        if(!isStartCode){
            break;
        }

        nalu.StartCodeLen = startCodeLen;

        int endPos = -1;
        for(int i = 2; i < bufferLen; i++){
            startCodeLen = 0;
            isStartCode = CheckStartCode(startCodeLen, buf + i, bufferLen - i);
            if(isStartCode){
                endPos = i;
                break;
            }
        }

        if(endPos > 0){
            nalu.SetBuf(buffer,endPos);
            uint8_t* _buffer = (uint8_t*)malloc(bufferLen - endPos);

            memcpy(_buffer, buffer+endPos, bufferLen - endPos);

            if(buffer != nullptr){
                free(buffer);
                buffer = nullptr;
            }

            buffer = _buffer;
            bufferLen = bufferLen - endPos;

            return 0;
        }
        else{
            if(isEnd == true){
                nalu.SetBuf(buffer,bufferLen);
                if(buffer != nullptr){
                    free(buffer);
                    buffer = nullptr;
                }
                buffer = nullptr;
                bufferLen = 0;            
                return 0;
            }
            int readedLen = ReadFromFile();
            if(readedLen <= 0){
                isEnd = true;
            }
        }



    }
    return -1;

}